import json
import base64
import io
import os
import sys
from PIL import Image

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from log_config import logger
from config import settings

# 优化后的系统提示词
SYSTEM_PROMPT = """请识别图片中的数学题目，将其分类到以下类别之一：
- 2D平面几何
- 3D立体几何  
- 函数
- 其他

要求：
1. 准确提取题目文字内容
2. 根据题目涉及的数学知识点进行分类
3. 严格按照指定的JSON格式返回结果

返回格式：
{
    "题目": "完整的题目内容",
    "类型": "分类结果（2D平面几何/3D立体几何/函数/其他）"
}"""

def classify_math_problem(image_path=None, question_text=""):
    """
    分类数学题目
    
    :param image_path: 可选的图片文件路径
    :param question_text: 可选的题目文字描述
    :return: 分类结果
    """
    try:
        logger.info(f"classify_math_problem 被调用，图片路径: {image_path}, 文字长度: {len(question_text) if question_text else 0}")
        
        # 检查输入参数
        if not image_path and not question_text.strip():
            logger.error("输入参数检查失败：图片和文字都为空")
            raise ValueError("必须提供图片路径或题目文字，不能两者都为空")
        
        base64_image = None
        
        # 处理图片（如果提供了图片路径）
        if image_path:
            logger.info(f"开始处理图片: {image_path}")
            
            # 检查文件是否存在
            if not os.path.exists(image_path):
                raise FileNotFoundError(f"图片文件不存在: {image_path}")
            
            # 读取并转换图片为base64
            with open(image_path, 'rb') as f:
                image_data = f.read()
            
            image = Image.open(io.BytesIO(image_data))
            
            # 转换为PNG格式的base64
            buffered = io.BytesIO()
            if image.mode in ('RGBA', 'LA', 'P'):
                image = image.convert('RGB')
            image.save(buffered, format="PNG")
            img_data = buffered.getvalue()
            base64_image = base64.b64encode(img_data).decode('ascii')
            
            logger.info("图片处理完成")
        
        logger.info("开始调用API...")
        logger.info(f"API配置 - base_url: {settings.version_base_url}, model: {settings.version_model}")
        
        # 调用API - 使用DeepSeek API替代无法访问的内网API
        from openai import OpenAI
        client = OpenAI(
            api_key=settings.deepseek_api_key, 
            base_url=settings.deepseek_base_url
        )
        
        # 构建消息内容
        user_content = []
        
        # 添加文字内容（如果提供了）
        if question_text.strip():
            user_content.append({
                "type": "text",
                "text": question_text
            })
        
        # 添加图片内容（如果提供了）
        if base64_image:
            user_content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{base64_image}",
                },
            })
        
        # 如果只有文字没有图片，直接使用文字内容
        if not base64_image:
            user_content = question_text
        
        logger.info(f"准备发送API请求，用户内容类型: {type(user_content)}")
        
        completion = client.chat.completions.create(
            model=settings.deepseek_chat_model,
            messages=[
                {
                    "role": "system",
                    "content": SYSTEM_PROMPT
                },
                {
                    "role": "user", 
                    "content": user_content
                }
            ]
        )
        
        response_text = completion.choices[0].message.content
        logger.info("API调用成功")
        
        # 解析JSON结果
        try:
            # 尝试直接解析JSON
            result = json.loads(response_text)
        except json.JSONDecodeError:
            # 如果失败，尝试提取JSON部分
            import re
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                result = json.loads(json_match.group())
            else:
                raise ValueError("无法解析API返回的JSON格式")
        
        # 验证返回格式
        if "题目" not in result or "类型" not in result:
            raise ValueError("API返回的JSON格式不正确")
        
        # 验证类型分类
        valid_types = ["2D平面几何", "3D立体几何", "函数", "其他"]
        if result["类型"] not in valid_types:
            logger.warning(f"分类结果 '{result['类型']}' 不在预期范围内，设置为'其他'")
            result["类型"] = "其他"
        
        logger.info(f"分类结果: {result}")
        
        # 保存结果
        output_dir = "api_test/output"
        os.makedirs(output_dir, exist_ok=True)
        
        final_result = {
            "success": True,
            "result": result,
            "raw_response": response_text
        }
        
        with open(f"{output_dir}/classification_result.json", 'w', encoding='utf-8') as f:
            json.dump(final_result, f, ensure_ascii=False, indent=4)
        
        logger.info(f"结果已保存到 {output_dir}/classification_result.json")
        
        return final_result
        
    except Exception as e:
        logger.error(f"处理失败: {e}")
        return {
            "success": False,
            "error": str(e)
        }

# 测试代码
if __name__ == '__main__':
    print("=" * 50)
    print("数学题目分类识别")
    print("=" * 50)
    
    # 测试场景选择
    test_mode = 3  # 1: 仅图片, 2: 仅文字, 3: 图片+文字
    
    if test_mode == 1:
        # 测试场景1: 仅使用图片
        print("🖼️ 测试模式: 仅图片识别")
        test_image_path = r"D:\AI教育\题目模型库\黄哥\question_cut\question_cut\image\2025gaokao1\cut_result\1_6.jpg"
        result = classify_math_problem(image_path=test_image_path)
        
    elif test_mode == 2:
        # 测试场景2: 仅使用文字
        print("📝 测试模式: 仅文字识别")
        test_question_text = "已知函数f(x) = x² + 2x + 1，求函数的最小值点"
        result = classify_math_problem(question_text=test_question_text)
        
    elif test_mode == 3:
        # 测试场景3: 图片+文字
        print("🖼️📝 测试模式: 图片+文字识别")
        test_image_path = r"D:\AI教育\题目模型库\黄哥\question_cut\question_cut\image\2025gaokao1\cut_result\1_6.jpg"
        test_question_text = "这是一道关于函数的题目"
        result = classify_math_problem(image_path=test_image_path, question_text=test_question_text)
    
    # 显示结果
    print("\n分类结果:")
    print("-" * 30)
    
    if result['success']:
        print("✅ 识别成功!")
        classification = result['result']
        print(f"题目: {classification['题目']}")
        print(f"类型: {classification['类型']}")
        print(f"📁 详细结果已保存到: api_test/output/classification_result.json")
    else:
        print("❌ 识别失败!")
        print(f"错误: {result.get('error', '未知错误')}")
    
    print("\n💡 提示: 修改 test_mode 变量来测试不同场景")
    print("   1: 仅图片识别")
    print("   2: 仅文字识别") 
    print("   3: 图片+文字识别")
    print("=" * 50)